import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras import layers, models
from utils.evalute_model_plot_metric import plot_metric


def load_data():
    """
    数据加载
    :return:
    """
    train_data = '../data/titanic/train.csv'
    test_data = '../data/titanic/test.csv'
    df_train = pd.read_csv(train_data)
    df_test = pd.read_csv(test_data)
    return df_train, df_test


class ExploratoryData(object):
    """
    探索性数据分析
    """
    def __init__(self, df_train):
        self.df_train = df_train

    def labels(self):
        ax = self.df_train['Survived'].value_counts().plot(kind='bar', rot=0)
        ax.set_ylabel('Counts')
        ax.set_xlabel('Survived')
        plt.show()

    def ages(self):
        ax = self.df_train['Age'].plot(kind='hist', bins=20)
        ax.set_ylabel('Freq')
        ax.set_xlabel('Age')
        plt.show()

    def corr_ages_labels(self):
        ax = self.df_train.query('Survived==0')['Age'].plot(kind='density')
        self.df_train.query('Survived==1')['Age'].plot(kind='density')
        ax.legend(['Survived = 0', 'Survived = 1'])
        ax.set_ylabel('Density')
        ax.set_xlabel('Age')
        plt.show()


class TfModel(object):
    """
    tf
    """
    def __init__(self, df_train, df_test):
        self.df_train = df_train
        self.df_test = df_test

    @staticmethod
    def preprocess(df):
        """
        预处理
        :param df:
        :return:
        """
        df_result = pd.DataFrame()

        # Pclass 乘客所持票类，有三种值(1,2,3)  get_dummies 转换成one hot encode
        df_pclass = pd.get_dummies(df['Pclass'])
        df_pclass.columns = ['Pclass_' + str(x) for x in df_pclass.columns]
        df_result = pd.concat([df_result, df_pclass], axis=1)

        # sex
        df_sex = pd.get_dummies(df['Sex'])
        df_result = pd.concat([df_result, df_sex], axis=1)

        # age
        df_result['Age'] = df['Age'].fillna(0)
        # 判断age是否为空
        df_result['Age_null'] = pd.isna(df['Age']).astype('int32')

        df_result['SibSp'] = df['SibSp']
        df_result['Parch'] = df['Parch']
        df_result['Fare'] = df['Fare']

        # 判断Cabin是否为空
        df_result['Cabin_null'] = pd.isna(df['Cabin']).astype('int32')

        # Embarked
        df_embark = pd.get_dummies(df['Embarked'], dummy_na=True)
        df_embark.columns = ['Embarked_' + str(x) for x in df_embark.columns]
        df_result = pd.concat([df_result, df_embark], axis=1)

        return df_result

    @staticmethod
    def tf_model(input_shape):
        """
        构建模型
        :param input_shape:
        :return:
        """
        # 清空会话
        tf.keras.backend.clear_session()

        # 最简单的Sequential，按层顺序模型
        model = models.Sequential()
        model.add(layers.Dense(20, activation='relu', input_shape=(input_shape,)))
        model.add(layers.Dense(10, activation='relu'))
        model.add(layers.Dense(1, activation='sigmoid'))

        return model

    @staticmethod
    def train_model(model, x_train, y_train):
        """
        模型训练: 二分类问题选择二元交叉熵损失函数: https://www.cnblogs.com/sunrise-keephungary/p/10056027.html
        :param model:
        :param x_train:
        :param y_train:
        :return:
        """
        model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['AUC'])
        # 交叉检验
        return model.fit(x_train, y_train, batch_size=64, epochs=30, validation_split=0.2)

    @staticmethod
    def evalute_model(model, history, x_test, y_test):
        """
        模型评估
        :param model:
        :param history:
        :param x_test:
        :param y_test:
        :return:
        """
        # loss
        plot_metric(history=history, metric='loss')
        # auc
        plot_metric(history=history, metric='auc')
        # scores
        scores = model.evaluate(x=x_test, y=y_test)
        print(scores)

    def main(self):
        # preprocess
        x_train = self.preprocess(df=self.df_train)
        y_train = self.df_train['Survived'].values

        x_test = self.preprocess(df=self.df_test)
        y_test = self.df_test['Survived'].values
        # 构建模型
        model = self.tf_model(input_shape=len(x_train.columns))
        # 训练模型
        history = self.train_model(model=model, x_train=x_train, y_train=y_train)
        # print(history.history)
        # 评估模型
        self.evalute_model(model=model, history=history, x_test=x_test, y_test=y_test)
        # print(model.summary())

        # predict
        # 概率预测
        pre_probability = model.predict(x_test[: 10])
        # 类别预测
        pre_class = model.predict_classes(x_test[: 10])
        print(pre_probability)
        print(pre_class)

        # model save and load
        # model_path = '../data/models/tf_model_titanic'
        # model.save(model_path)
        # model_load = tf.keras.models.load_model(model_path)


def run():
    df_train, df_test = load_data()
    """
    # EDA
    explore = ExploratoryData(df_train=df_train)
    explore.labels()
    explore.ages()
    explore.corr_ages_labels()
    """
    tf_model = TfModel(df_train=df_train, df_test=df_test)
    tf_model.main()


if __name__ == '__main__':
    run()
